package utils

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
)

// GenRsaKey 生成私钥和公钥
func GenRsaKey() (priKey, pubKey []byte) {
	// 生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		panic(err)
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	// 生成私钥
	priKey = pem.EncodeToMemory(block)

	publicKey := &privateKey.PublicKey
	derPkix := x509.MarshalPKCS1PublicKey(publicKey)
	if err != nil {
		panic(err)
	}
	block = &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: derPkix,
	}
	// 生成公钥
	pubKey = pem.EncodeToMemory(block)
	return
}

// ParsePrivateKeyBytes pkcs1私钥解析: 解密pem格式的私钥
func ParsePrivateKeyBytes(privateKey []byte) (*rsa.PrivateKey, error) {
	p := &pem.Block{}
	p, _ = pem.Decode(privateKey)
	if p == nil {
		return nil, errors.New("ParsePrivateKeyBytes pem.Decode failed")
	}
	return x509.ParsePKCS1PrivateKey(p.Bytes)
}

// ParsePubKeyBytes pkcs1公钥解析: : 解密pem格式的公钥
func ParsePubKeyBytes(publicKey []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(publicKey)
	if block == nil {
		return nil, errors.New("ParsePubKeyBytes pem.Decode failed")
	}
	return x509.ParsePKCS1PublicKey(block.Bytes)
}

// RsaSignWithSha256 签名
func RsaSignWithSha256(data []byte, priKey []byte) []byte {
	h := sha256.New()
	h.Write(data)
	hashed := h.Sum(nil)
	block, _ := pem.Decode(priKey)
	if block == nil {
		panic(errors.New("private key error"))
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		fmt.Println("ParsePKCS8PrivateKey err", err)
		panic(err)
	}

	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hashed)
	if err != nil {
		fmt.Printf("Error from signing: %s\n", err)
		panic(err)
	}

	return signature
}

// RsaVerifyWithSha256 验证
func RsaVerifyWithSha256(data, signData, keyBytes []byte) bool {
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		panic(errors.New("public key error"))
	}
	pubKey, err := x509.ParsePKCS1PublicKey(block.Bytes)
	if err != nil {
		panic(err)
	}

	hashed := sha256.Sum256(data)
	err = rsa.VerifyPKCS1v15(pubKey, crypto.SHA256, hashed[:], signData)
	if err != nil {
		panic(err)
	}
	return true
}

// RsaEncryptByPublicKey 公钥加密
func RsaEncryptByPublicKey(data, publicKey []byte) []byte {
	// 解析公钥
	pub, err := ParsePubKeyBytes(publicKey)
	if err != nil {
		panic(err)
	}
	//加密
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, pub, data)
	if err != nil {
		panic(err)
	}
	return ciphertext
}

// RsaDecryptByPrivateKey 私钥解密
func RsaDecryptByPrivateKey(ciphertext, privteKey []byte) []byte {
	// 解析私钥
	priv, err := ParsePrivateKeyBytes(privteKey)
	if err != nil {
		panic(err)
	}
	// 解密
	data, err := rsa.DecryptPKCS1v15(rand.Reader, priv, ciphertext)
	if err != nil {
		panic(err)
	}
	return data
}
